package com.base.data.elasticsearch.util;

import com.google.common.collect.Lists;
import lombok.extern.slf4j.Slf4j;
import org.elasticsearch.action.bulk.BulkRequest;
import org.elasticsearch.action.delete.DeleteRequest;
import org.elasticsearch.action.support.ActiveShardCount;
import org.elasticsearch.action.support.WriteRequest;
import org.elasticsearch.action.update.UpdateRequest;
import org.elasticsearch.common.unit.TimeValue;
import org.springframework.util.CollectionUtils;

import java.util.Iterator;
import java.util.List;
import java.util.function.Function;
import java.util.function.Supplier;
import java.util.stream.Collectors;

/**
 * Bulk 默认设置批量提交的数据量不能超过 100M,超过：413 Request Entity Too Large
 *  request : request.sizeInBytes/1024/1024
 */
@Slf4j
public class BulkRequestSpitUtil {

    private static final int maxBytes = 75 * 1024 * 1024;//默认给75m

    private static BulkRequest defaultBulkReqConfig(BulkRequest param) {
        BulkRequest request = param == null ? new BulkRequest() : param;
        request.timeout(TimeValue.timeValueMinutes(3));
        request.setRefreshPolicy(WriteRequest.RefreshPolicy.WAIT_UNTIL);
        request.waitForActiveShards(ActiveShardCount.DEFAULT);
        return request;
    }

    public static <T> List<BulkRequest> spitUpdateRequest(Supplier<List<T>> supplier, Function<T, UpdateRequest> function) {
        List<BulkRequest> requests = Lists.newArrayList();
        final List<T> queries = supplier.get();
        if (!CollectionUtils.isEmpty(queries)) {
            log.debug("{} - [{}]", Thread.currentThread().getStackTrace()[1].getMethodName(),queries.size());
            BulkRequest request = defaultBulkReqConfig(null);
            final Iterator<T> iterator = queries.iterator();
            while (iterator.hasNext()) {
                final T indexQuery = iterator.next();
                request.add(function.apply(indexQuery));
                if (request.estimatedSizeInBytes() > 0 && request.estimatedSizeInBytes() > maxBytes) {
                    requests.add(request);
                    request = defaultBulkReqConfig(null);
                }
                if (request.estimatedSizeInBytes() > 0 && !iterator.hasNext()) {
                    requests.add(request);
                }
            }
        }
        log.info("{} - [{}]", Thread.currentThread().getStackTrace()[1].getMethodName(),requests.stream().map(p -> String.valueOf(p.numberOfActions())).collect(Collectors.joining(";")));
        return requests;
    }

    public static List<BulkRequest> spitDeleteRequest(List<String> queries, String indexName,String typeName) {
        List<BulkRequest> requests = Lists.newArrayList();
        if (!CollectionUtils.isEmpty(queries)) {
            log.debug("{} - [{}]", Thread.currentThread().getStackTrace()[1].getMethodName(),queries.size());
            BulkRequest request = defaultBulkReqConfig(null);
            final Iterator<String> iterator = queries.iterator();
            while (iterator.hasNext()) {
                final String item = iterator.next();
                request.add(new DeleteRequest(indexName, typeName, item));
                if (request.estimatedSizeInBytes() > 0 && request.estimatedSizeInBytes() > maxBytes) {
                    requests.add(request);
                    request = defaultBulkReqConfig(null);
                }
                if (request.estimatedSizeInBytes() > 0 && !iterator.hasNext()) {
                    requests.add(request);
                }
            }
        }
        log.info("{} - [{}]", Thread.currentThread().getStackTrace()[1].getMethodName(),requests.stream().map(p -> String.valueOf(p.numberOfActions())).collect(Collectors.joining(";")));
        return requests;
    }

    public static void initParamUpdateRequest(UpdateRequest request) {
        request.docAsUpsert(true);
        request.timeout(TimeValue.timeValueMinutes(2));
        request.setRefreshPolicy(WriteRequest.RefreshPolicy.WAIT_UNTIL);
        request.waitForActiveShards(ActiveShardCount.DEFAULT);
        request.retryOnConflict(20);
    }

    /**
     * refreshPolicy= {@link WriteRequest.RefreshPolicy#NONE}
     * to fixed: RefreshPolicy is not supported on an item request. Set it on the BulkRequest instead.
     * @param request
     * @param refreshPolicy
     */
    public static void initParamUpdateRequest(UpdateRequest request, WriteRequest.RefreshPolicy refreshPolicy) {
        request.docAsUpsert(true);
        request.timeout(TimeValue.timeValueMinutes(2));
        request.setRefreshPolicy(refreshPolicy);
        request.waitForActiveShards(ActiveShardCount.DEFAULT);
        request.retryOnConflict(20);
    }
}
